import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import cv2
import numpy as np
from tensorflow.keras.models import load_model
import tensorflow.keras.backend as K
import tensorflow as tf
from Tool.Estimate import Estimate
from tensorflow import keras
tf.compat.v1.disable_eager_execution()

def _grad_cam_image(image,model_input,cam_conv_tensor,class_out_tensor):
    src_img = image.copy()
    h,w = image.shape[:2]
    img = image / 255.
    image_dim =  np.expand_dims(img,axis=0)

    grad = K.gradients(class_out_tensor,cam_conv_tensor)[0]

    weight = K.mean(grad,axis=(0,1,2))
    iter = K.function([model_input],[weight,cam_conv_tensor[0]])

    weight_value,cam_conv_value = iter(image_dim)
    depth = weight_value.shape[-1]
    for i in range(depth):
        cam_conv_value[:,:,i] *= weight_value[i]

    headmap = np.mean(cam_conv_value,axis=-1)
    headmap = np.maximum(headmap,0)
    headmap /= np.max(headmap)

    headmap = np.uint8(255*cv2.resize(headmap,(w,h)))
    headmap = cv2.applyColorMap(headmap, cv2.COLORMAP_JET)
    org_headmap = headmap.copy()

    heat_image = cv2.addWeighted(src_img, 0.6, headmap, 0.4, 0)
    return heat_image,org_headmap


if __name__ == '__main__':
    model = load_model('../out/grad_tap_9_001-0.3750',custom_objects={"Estimate":Estimate})
    # classer = Classify()
    img = cv2.imread('../test_data/1/201912230001_20200421001.jpg')
    image = cv2.resize(img, (512, 512))
    image_src = image.copy()
    # orc_img,heat_map,org_heat_map = classer.grad_cam_image(img)
    cam_conv_tensor = model.get_layer(index=11).output
    class_id = 0
    class_out_tensor = model.output[:, class_id]
    heat_map, org_heat_map = _grad_cam_image(image, model.input, cam_conv_tensor, class_out_tensor)
    cv2.imwrite('../out/orc_img.jpeg', image_src)
    cv2.imwrite('../out/heat_map.jpeg', heat_map)
    cv2.imwrite('../out/org_heat_map.jpeg', org_heat_map)
